from urllib.parse import urlparse
import base64
import hashlib
import hmac
import json
from datetime import datetime
from time import mktime
from urllib.parse import urlencode
from wsgiref.handlers import format_date_time
import websocket
import ssl
from spark_config import SPARK_CONFIG, WS_URL

# 从配置文件读取参数
appid = SPARK_CONFIG['APPID']
api_secret = SPARK_CONFIG['APISecret']
api_key = SPARK_CONFIG['APIKey']
imageunderstanding_url = WS_URL

class Ws_Param(object):
    """
    WebSocket参数类
    """
    def __init__(self, APPID, APIKey, APISecret, imageunderstanding_url):
        """
        初始化
        """
        self.APPID = APPID
        self.APIKey = APIKey
        self.APISecret = APISecret
        self.host = urlparse(imageunderstanding_url).netloc
        self.path = urlparse(imageunderstanding_url).path
        
    def create_url(self):
        """
        生成url
        """
        # 生成RFC1123格式的时间戳
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 拼接字符串
        signature_origin = "host: " + self.host + "\n"
        signature_origin += "date: " + date + "\n"
        signature_origin += "GET " + self.path + " HTTP/1.1"

        # 进行hmac-sha256进行加密
        signature_sha = hmac.new(self.APISecret.encode('utf-8'), signature_origin.encode('utf-8'),
                               digestmod=hashlib.sha256).digest()

        signature_sha_base64 = base64.b64encode(signature_sha).decode(encoding='utf-8')

        authorization_origin = f'api_key="{self.APIKey}", algorithm="hmac-sha256", headers="host date request-line", signature="{signature_sha_base64}"'

        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        # 将请求的鉴权参数组合为字典
        v = {
            "authorization": authorization,
            "date": date,
            "host": self.host
        }
        # 拼接鉴权参数，生成url
        url = imageunderstanding_url + '?' + urlencode(v)
        return url

# 全局变量用于存储回答
answer = ""

# 收到websocket错误的处理
def on_error(ws, error):
    """
    收到websocket错误的处理
    """
    print("### error:", error)

# 收到websocket关闭的处理
def on_close(ws, close_status_code, close_msg):
    """
    收到websocket关闭的处理
    """
    print("### closed ###")

# 收到websocket连接建立的处理
def on_open(ws):
    """
    收到websocket连接建立的处理
    """
    def run(*args):
        data = json.dumps(gen_params(ws.appid, ws.question))
        ws.send(data)

    run()

# 收到websocket消息的处理
def on_message(ws, message):
    """
    收到websocket消息的处理
    """
    global answer
    resp_data = json.loads(message)
    
    if resp_data['header']['code'] != 0:
        print(f'请求错误: {resp_data["header"]}')
        ws.close()
    else:
        choices = resp_data["payload"]["choices"]
        status = choices["status"]
        content = choices["text"][0]["content"]
        answer += content
        
        # 当响应结束时关闭连接
        if status == 2:
            ws.close()

def gen_params(appid, question):
    """
    通过appid和用户的提问来生成请参数
    """
    # 如果传入的是列表（对话上下文），直接使用
    messages = question if isinstance(question, list) else [{"role": "user", "content": question}]
    
    data = {
        "header": {
            "app_id": appid,
            "uid": "1234"
        },
        "parameter": {
            "chat": {
                "domain": "general",
                "temperature": 0.5,
                "max_tokens": 1024
            }
        },
        "payload": {
            "message": {
                "text": messages
            }
        }
    }
    return data

# 全局变量，用于存储当前的对话上下文
current_context = []

def main(appid, api_key, api_secret, imageunderstanding_url, question):
    """
    主函数：处理图片理解请求并返回结果
    """
    global answer, current_context
    answer = ""  # 重置全局变量
    current_context = question  # 保存对话上下文，包含图片数据
    
    wsParam = Ws_Param(appid, api_key, api_secret, imageunderstanding_url)
    websocket.enableTrace(False)
    wsUrl = wsParam.create_url()
    ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open)
    ws.appid = appid
    ws.question = question
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
    
    # 记录助手的回答到上下文
    current_context = current_context + [{"role": "assistant", "content": answer}]
    
    return answer

def continue_dialog(appid, api_key, api_secret, imageunderstanding_url, follow_up_question):
    """
    继续图片相关的对话
    """
    global answer, current_context
    answer = ""  # 重置回答
    
    # 构建完整的对话上下文
    context = current_context + [{"role": "user", "content": follow_up_question}]
    
    wsParam = Ws_Param(appid, api_key, api_secret, imageunderstanding_url)
    websocket.enableTrace(False)
    wsUrl = wsParam.create_url()
    ws = websocket.WebSocketApp(wsUrl, on_message=on_message, on_error=on_error, on_close=on_close, on_open=on_open)
    ws.appid = appid
    ws.question = context
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
    
    # 更新上下文，添加新的回答
    current_context.extend([
        {"role": "user", "content": follow_up_question},
        {"role": "assistant", "content": answer}
    ])
    
    return answer

def getText(role, content, content_type="text"):
    """
    生成文本或图片消息
    """
    if content_type == "image":
        return {"role": role, "content": content, "content_type": "image"}
    return {"role": role, "content": content}

def getlength(text):
    """
    获取字符串长度
    """
    length = 0
    for content in text:
        temp = content["content"]
        if type(temp) == str:
            length += len(temp)
    return length

def checklen(text):
    """
    检查文本长度是否超限
    """
    return getlength(text) <= 8192

if __name__ == '__main__':

    #text.clear
    while(1):
        Input = input("\n" +"问:")
        question = checklen(getText("user",Input))
        answer = ""
        print("答:",end = "")
        main(appid, api_key, api_secret, imageunderstanding_url, question)
        getText("assistant", answer)
        # print(str(text))

